"""
@Filename       : dataset_evaluation.py
@Create Time    : 2020/11/12 10:33
@Author         : Rylynn
@Description    : 

"""
import os


def count_user(dataset_path):
    user_set = []
    sequence_count = 0
    sequence_length_count = 0
    with open(os.path.join(dataset_path, 'train.txt')) as f:
        for line in f.readlines():
            query, cascade = line.split(' ', maxsplit=1)
            users = cascade.split(' ')[::2]
            sequence_count = sequence_count + 1
            sequence_length_count = sequence_length_count + 1 + len(users)
            for user in users:
                user = user.replace('\n', '')
                if user not in user_set:
                    user_set.append(user)

    with open(os.path.join(dataset_path, 'test.txt')) as f:
        for line in f.readlines():
            query, cascade = line.split(' ', maxsplit=1)
            users = cascade.split(' ')[::2]
            sequence_count = sequence_count + 1
            sequence_length_count = sequence_length_count + 1 + len(users)
            for user in users:
                user = user.replace('\n', '')
                if user not in user_set:
                    user_set.append(user)

    with open(os.path.join(dataset_path, 'graph.txt')) as f:
        for line in f.readlines():
            source, target = line.split(' ')
            source = source.replace('\n', '')
            target = target.replace('\n', '')
            if source not in user_set:
                user_set.append(source)
            if target not in user_set:
                user_set.append(target)


    print('user size:', len(user_set))
    print('action count', sequence_length_count)
    print('sequence count:', sequence_count)
    print('average length:', float(sequence_length_count) / float(sequence_count))

count_user('../../data/memes')